{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c7a36f33",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'S-46.62=-15.67-12.66+-18.29E'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%run common.ipynb\n",
    "\n",
    "tokenizer.decode(tokenizer.get_data(third_number=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f646376b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/pt2/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ModelGEN(\n",
       "  (feature): LlamaModel(\n",
       "    (embed_tokens): Embedding(22, 64)\n",
       "    (layers): ModuleList(\n",
       "      (0-3): 4 x LlamaDecoderLayer(\n",
       "        (self_attn): LlamaAttention(\n",
       "          (q_proj): Linear(in_features=64, out_features=64, bias=False)\n",
       "          (k_proj): Linear(in_features=64, out_features=64, bias=False)\n",
       "          (v_proj): Linear(in_features=64, out_features=64, bias=False)\n",
       "          (o_proj): Linear(in_features=64, out_features=64, bias=False)\n",
       "          (rotary_emb): LlamaRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): LlamaMLP(\n",
       "          (gate_proj): Linear(in_features=64, out_features=64, bias=False)\n",
       "          (up_proj): Linear(in_features=64, out_features=64, bias=False)\n",
       "          (down_proj): Linear(in_features=64, out_features=64, bias=False)\n",
       "          (act_fn): SiLUActivation()\n",
       "        )\n",
       "        (input_layernorm): LlamaRMSNorm()\n",
       "        (post_attention_layernorm): LlamaRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): LlamaRMSNorm()\n",
       "  )\n",
       "  (fc_out): Linear(in_features=64, out_features=22, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_dpo = torch.load('dpo.model')\n",
    "model_dpo.to(device)\n",
    "model_dpo.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "de9d94bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "S35.75= -3.68+39.08+0.52E 0.1700000000000017\n",
      "S129.94= 20.66--89.97+19.86E 0.5500000000000114\n",
      "S-77.23= 63.49/-14.79+-73.67E 0.7327653820148754\n",
      "S8346.47= -96.48*-87.66+-49.57E 61.39680000000044\n",
      "S-38.34= 39.43/97.51+-38.56E 0.1843687826889564\n",
      "S102.18= 36.68-9.98+76.29E 0.8100000000000023\n",
      "S-66.36= -5.25-64.38+6.08E 2.8100000000000023\n",
      "S116.83= 61.96+-21.66+78.47E 1.9399999999999977\n",
      "S-95.93= -14.80+12.64+-94.56E 0.789999999999992\n",
      "S-177.68= 6.56-98.00+-83.68E 2.5600000000000023\n",
      "S2198.58= -39.34*-57.20+-2.46E 49.20800000000054\n",
      "S3384.91= 79.58*42.10+-21.93E 56.52199999999948\n",
      "S2333.02= 69.64*33.99+-83.21E 49.16640000000007\n",
      "S1149.11= -71.22*-14.15+85.16E 56.1869999999999\n",
      "S63.19= 23.36+34.59+5.27E 0.030000000000001137\n",
      "S98.02= 87.57+-16.33+26.06E 0.7199999999999989\n",
      "S90.55= -21.48/45.30+90.15E 0.8741721854304529\n",
      "S127.17= 84.55+-48.06+91.37E 0.6899999999999977\n",
      "S-55.52= -61.50/-37.93+-57.72E 0.5785921434220924\n",
      "S-13.52= 43.77-83.07+26.76E 0.9800000000000111\n",
      "S-189.71= -88.53+-58.18+-43.45E 0.45000000000001705\n",
      "S-25.64= 82.81+-13.52+-96.04E 1.1099999999999994\n",
      "S69.53= -87.69+64.34+91.01E 1.8699999999999903\n",
      "S31.91= -17.78-35.89+84.53E 1.0500000000000007\n",
      "S-953.13= -13.88*63.79+-30.99E 36.73479999999995\n",
      "S-84.39= 84.98/-85.44+-83.59E 0.19461610486891345\n",
      "S-48.52= -47.74+-88.30+87.28E 0.23999999999998778\n",
      "S-75.90= 55.69-93.95+-39.37E 1.7299999999999898\n",
      "S4664.86= -89.45*-51.58+20.83E 30.198999999999614\n",
      "S-76.63= -38.19-34.83+-4.91E 1.2999999999999972\n",
      "S-234.82= -65.01+-89.44+-84.46E 4.089999999999975\n",
      "S19.91= 45.12+25.58+-50.07E 0.7199999999999882\n",
      "S94.39= 81.13+91.50+-79.62E 1.3800000000000097\n",
      "S31.36= -26.78/44.10+31.44E 0.5272562358276609\n",
      "S38.99= 78.00-34.23+-6.28E 1.5\n",
      "S-136.35= -13.63-26.75+-93.48E 2.4899999999999807\n",
      "S-54.76= 21.15/41.84+-55.80E 0.5345028680688344\n",
      "S106.71= 41.30--64.59+0.67E 0.14999999999999147\n",
      "S-217.96= -89.33+-67.39+-64.23E 2.9899999999999807\n",
      "S105.19= 16.50--32.88+56.32E 0.5100000000000051\n",
      "S69.50= 68.91/-16.51+73.76E 0.08616596002423194\n",
      "S43.24= -95.92/90.20+44.51E 0.2065853658536554\n",
      "S-71.69= 15.38/47.88+-71.94E 0.07121971595655907\n",
      "S-41.58= -8.90-9.88+-20.34E 2.4599999999999937\n",
      "S-90.11= -9.46-19.28+-58.55E 2.8200000000000074\n",
      "S-31.74= -50.31--72.41+-53.06E 0.7799999999999905\n",
      "S67.83= 83.26-26.20+10.62E 0.14999999999999147\n",
      "S42.86= 94.90/9.11+32.59E 0.14712403951701702\n",
      "S-149.07= -50.42+-38.08+-60.06E 0.5099999999999909\n",
      "S-97.50= 99.22/-84.25+-96.43E 0.10768545994065448\n",
      "S-67.80= -60.83--28.90+-34.02E 1.8499999999999943\n",
      "S-82.50= 79.45/-68.61+-81.19E 0.15200553855123644\n",
      "S-479.88= 55.22*-5.87+-92.14E 63.59860000000003\n",
      "S39.64= 25.87/-27.74+40.67E 0.09741167988464383\n",
      "S-191.47= 49.81*-4.65+36.30E 3.8465000000000202\n",
      "S56.33= 89.61--72.04+-99.09E 6.230000000000004\n",
      "S34.48= 71.64+22.65+-57.48E 2.3299999999999983\n",
      "S-4808.33= 46.09*-99.59+-92.72E 125.5068999999994\n",
      "S17.23= -89.84/-59.24+15.52E 0.19345712356516032\n",
      "S-109.56= -71.98+65.09+-99.22E 3.450000000000003\n",
      "S42.97= -6.48+21.50+27.68E 0.269999999999996\n",
      "S-38.96= 32.03/-73.17+-38.51E 0.012252289189561338\n",
      "S-100.07= -92.59+79.79+-86.76E 0.5099999999999909\n",
      "S1349.37= 82.86*16.56+-39.41E 16.618400000000065\n"
     ]
    }
   ],
   "source": [
    "#随机一批数据\n",
    "input_ids = [tokenizer.get_data(third_number=True) for i in range(64)]\n",
    "\n",
    "#切分成question和answer\n",
    "split = [i.index(tokenizer.encoder['=']) + 1 for i in input_ids]\n",
    "question = [input_ids[i][:split[i]] for i in range(len(input_ids))]\n",
    "answer = [input_ids[i][split[i]:] for i in range(len(input_ids))]\n",
    "\n",
    "#根据question生成predict\n",
    "input_ids = [torch.LongTensor(i).unsqueeze(0).to(device) for i in question]\n",
    "predict = [generate(model_dpo, i) for i in input_ids]\n",
    "\n",
    "#裁剪,只要生成的部分\n",
    "predict = [p[0].tolist()[len(q):] for p, q in zip(predict, question)]\n",
    "\n",
    "#解码成文本\n",
    "question = [tokenizer.decode(i) for i in question]\n",
    "answer = [tokenizer.decode(i) for i in answer]\n",
    "predict = [tokenizer.decode(i) for i in predict]\n",
    "\n",
    "for q, a, p in zip(question, answer, predict):\n",
    "    try:\n",
    "        diff = abs(float(q[1:-1]) - eval(p[:p.index('E')]))\n",
    "    except:\n",
    "        diff = abs(float(q[1:-1]))\n",
    "\n",
    "    print(q, p, diff)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt2]",
   "language": "python",
   "name": "conda-env-pt2-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
